#include "mpmdComm.hpp"
#include <iostream>
#include <cassert>
namespace EXTREAM {
MpmdComm *MpmdComm::instance = nullptr;

void MpmdComm::init(MPI_Comm peer_communicator) {
  MPI_Comm_dup(peer_communicator, &peer_comm_);

  MPI_Comm_size(peer_comm_, &peer_comm_size_);
  MPI_Comm_rank(peer_comm_, &rank_in_peer_);

  int *mpi_app_num, flag;
  MPI_Comm_get_attr(peer_comm_, MPI_APPNUM, &mpi_app_num, &flag);
  app_id_ = *mpi_app_num;

  MPI_Comm_split(peer_comm_, app_id_, rank_in_peer_, &local_comm_);

  MPI_Comm_size(local_comm_, &local_comm_size_);
  MPI_Comm_rank(local_comm_, &rank_in_local_);
  remote_comm_size_ = peer_comm_size_ - local_comm_size_;

  assert(app_id_ == 0 || app_id_ == 1);

  // local_leader 均为各个组 rank = 0的进程
  int leader_send[2]{-1, -1}, leader_recv[2]{-1, -1};
  if (rank_in_local_ == 0) {
    std::cout << "APPNUM: " << app_id_
              << "  leader_global_rank = " << rank_in_peer_ << std::endl;
    leader_send[app_id_] = rank_in_peer_;
  }
  MPI_Allreduce(leader_send, leader_recv, 2, MPI_INT, MPI_MAX, peer_comm_);

  local_leader_ = 0;
  if (app_id_ == 0)
    remote_leader_ = leader_recv[1];
  else
    remote_leader_ = leader_recv[0];

  MPI_Intercomm_create(local_comm_, local_leader_, peer_comm_, remote_leader_,
                       INTER_COMM_TAG, &inter_comm_);

  is_init_ = true;
}

void MpmdComm::initSPMD(MPI_Comm comm) {
  MPI_Comm_dup(comm, &local_comm_);
  MPI_Comm_rank(local_comm_, &rank_in_local_);
  MPI_Comm_size(local_comm_, &local_comm_size_);
  is_init_ = true;
  is_mpmd_ = false;
}

void MpmdComm::release() {
  delete instance;
  instance = nullptr;
}

MpmdComm::~MpmdComm() {
  MPI_Comm_free(&local_comm_);
  if (is_mpmd_) {
    MPI_Comm_free(&peer_comm_);
    MPI_Comm_free(&inter_comm_);
  }
}

} // namespace EXTREAM